package cs235.avl;

import java.util.Iterator;

public class AVLSetTree<T extends Comparable<T>> extends TreeSet<T>{
    private int size;
    private int balanceFactor = 2;
    
    private T value;
    private boolean done;
    
    private Node<T>[] up;
    private int[] upd;

    @SuppressWarnings("hiding")
    private class Node<T extends Comparable<T>> implements BinaryTreeNode {
        private T value;
        private Node<T>[] children;
        private int height;
        

        @SuppressWarnings("unchecked")
        public Node(T value) {
            this.value = value;
            children = new Node[2];
        }

        @Override
        public Object getData() {
            return value;
        }

        @Override
        public BinaryTreeNode getLeftChild() {
            return children[0];
        }

        @Override
        public BinaryTreeNode getRightChild() {
            return children[1];
        }

        @Override
        public int getHeight() {
            return height;
        }
    }

    private Node<T> root;

    private int not(int dir) {
        if (dir == 1) {
            return 0;
        }
        return 1;
    }

    private Node<T> singleRotation(Node<T> node, int dir) {
        Node<T> temp = node.children[not(dir)];
        
        int rlh;
        int rrh;
        int tlh;
        
        node.children[not(dir)] = temp.children[dir];
        temp.children[dir] = node;
        
        rlh = height( node.children[0] );
        rrh = height( node.children[1] );
        tlh = height( temp.children[not(dir)] );
        
        node.height = max(rlh, rrh) + 1;
        temp.height = max(tlh, node.height) + 1;
        
        return temp;
    }

    private Node<T> doubleRotation(Node<T> node, int dir) {
        node.children[not(dir)] = singleRotation(node.children[not(dir)], not(dir));
        return singleRotation(node, dir);
    }


    public Node<T> insert(Node<T> node) {
        
        if (node == null) {
            node = new Node<T>(value);
            size++;
        } else if (node.value.equals(this.value)){
            return node;
        }else {
            int dir = getDir(node.value, value);
            
            node.children[dir] = insert(node.children[dir]);
            
            if (!done){
                node = insertRebalance(node, dir);
            } 
        }
        return node;
    }

    private Node<T> insertRebalance(Node<T> node, int dir) {
        int lh;
        int rh;
        int max;
        lh = height(node.children[dir]);
        rh = height(node.children[not(dir)]);
        
        if (lh - rh >= balanceFactor){
            Node<T> a = node.children[dir].children[dir];
            Node<T> b = node.children[dir].children[not(dir)];
            
            if (height(a) >= height(b)){
                node = singleRotation(node, not(dir));
            } else {
                node = doubleRotation(node, not(dir));
            }
            
            done = true;
        }
        
        lh = height(node.children[dir]);
        rh = height(node.children[not(dir)]);
        max = max(lh, rh);
        
        node.height = max + 1;
        return node;
    }
    
    private boolean insert(T value){
        int start = size;
        done = false;
        this.value = value;
        root = insert(root);
        return !(size == start);
    }


    private int height(Node<T> node){
        return node == null ? -1 : node.height;
    }
    
    private int max(int a, int b){
        return a > b ? a : b;
    }
    


    private int getDir(T valueA, T valueB) {
        if (valueA.compareTo(valueB) < 0) {
            return 1;
        }
        return 0;
    }


    private int boolToC(boolean trueFalse) {
        if (trueFalse) {
            return 1;
        }
        return 0;
    }
    
    

    @SuppressWarnings("unchecked")
    public boolean remove(T value) {
        if (root != null) {
            up = new Node[size+1];
            upd = new int[size+1];
            Node<T> temp;
            int top = 0;
            
            temp = root;
            
            while(true){
                if (temp == null){
                    return false;
                }
                else if ( temp.value.equals(value)){
                    break;
                }
                
                upd[top] = getDir(temp.value, value);
                up[top++] = temp;
                
                temp = temp.children[upd[top-1]];
            }
            
            top = deleteThisOrChild(temp, top);
            
            while ( --top >= 0){
                if (balanceAfterDelete(top)){
                    break;
                }
            }
        } else {
            return false;
        }
        size--;
        return true;
    }

    private int deleteThisOrChild(Node<T> temp, int top) {
        if (temp.children[0] == null || temp.children[1] == null){
            int dir = boolToC(temp.children[0] == null);
            
            if (top != 0){
                up[top-1].children[upd[top-1]] = temp.children[dir];
            } else {
                root = temp.children[dir];
            }
            
            deleteNode(temp);
            
        } else {
            Node<T> heir = temp.children[1];
            
            
            upd[top] = 1;
            up[top++] = temp;
            
            while ( heir.children[0] != null){
                upd[top] = 0;
                up[top++] = heir;
                heir = heir.children[0];
            }
            
            temp.value = heir.value;
            
            up[top-1].children[boolToC(up[top-1] == temp)] = heir.children[1];
            
            deleteNode(heir);
        }
        return top;
    }

    private boolean balanceAfterDelete(int top) {
        int lh = height( up[top].children[upd[top]]);
        int rh = height( up[top].children[not(upd[top])]);
        int max = max(lh, rh);
        
        up[top].height = max +1;
        
        if ( lh - rh == -1 ){
            return true;
        }
        if ( lh - rh <= -balanceFactor ){
            Node<T> a = up[top].children[not(upd[top])].children[upd[top]];
            Node<T> b = up[top].children[not(upd[top])].children[not(upd[top])];
            
            if ( height(a) <= height(b)){
                up[top] = singleRotation(up[top], upd[top]);
            } else {
                up[top] = doubleRotation(up[top], upd[top]);
            }
            
            if ( top != 0 ){
                up[top-1].children[upd[top-1]] = up[top];
            } else {
                root = up[0];
            }
        }
        return false;
    }

    
    

    private void deleteNode(Node<T> node) {
        node.children = null;
        node.value = null;
    }

    @SuppressWarnings({
            "unchecked", "rawtypes"
    })
    @Override
    public Iterator iterator() {
        return new TreeSetIterator(toArray());
    }

    @Override
    public boolean add(T e) {
        return insert(e);
    }


    @SuppressWarnings("unchecked")
    @Override
    public boolean remove(Object o) {
        if (o instanceof Comparable<?>){
            return remove((T)o);
        }
        return false;
    }


    @Override
    public void clear() {
        clear(root);
        root = null;
        size = 0;
    }
    
    private void clear(Node<T> node){
        if (node != null){
            clear(node.children[0]);
            clear(node.children[1]);
            deleteNode(node);
        }
    }

    @Override
    public int size() {
        return size;
    }

    @SuppressWarnings("unchecked")
    @Override
    public boolean contains(Object o) {
        if (o instanceof Comparable<?>){
            Node<T> temp = root;
            while (temp != null){
                if (o.equals(temp.getData())){
                    return true;
                }
                temp = temp.children[getDir((T)temp.getData(), (T)o)];
            }
        }
        return false;
    }

    @Override
    public BinaryTreeNode getRootNode() {
        return root;
    }



}
